import sys
import os

# 添加 project/ 目录到 sys.path
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

import rosbag
import numpy as np
import sensor_msgs.point_cloud2 as pc2
import cv2
import tqdm
from cv_bridge import CvBridge
import csv
import json
import argparse
import multiprocessing
from utils import topic2path
from my_io import write_points, copy_file, files_filter
import time
from tictoc import TicToc
from cyw_devkit.core.zmath import quaternion_to_euler_angles

cls2dis = {
    'Car': 80,
    'Bus': 100,
    'Van': 100,
    'Pedestrian': 50,
    'Truck': 100,
    'Bicycle': 50,
    'Motorcycle': 60
}
cls2idx = {
    'Car': 6,
    'Bus': 2,
    'Van': 3,
    'Pedestrian': 4,
    'Truck': 7,
    'Bicycle': 5,
    'Motorcycle': 8
}
idx2cls = {val: key for key, val in cls2idx.items()}

now = time.time()


# 防止多进程时间戳重复，用进程id拉开一小时 3600*1000ms间隔
def get_ms(msg, process_id=0):
    if msg.header.stamp.secs == 0 and msg.header.stamp.nsecs == 0:
        return int(msg.time_us / 1000)
    return (msg.header.stamp.secs + int(now)) * 1000 + int(
        msg.header.stamp.nsecs / (10 ** 6) + process_id * 3600 * 1000)


def main(args):
    # 统计耗时
    cost = TicToc("ros解析")
    assert os.path.exists(args.rosbag_dir)
    files = files_filter(args.rosbag_dir, pos_fix="bag", split='.')
    process_size = len(files)
    manager = multiprocessing.Manager()
    if process_size > 1:
        pool = multiprocessing.Pool(process_size)
        counter_list = manager.list()
        for idx in range(process_size):
            pool.apply_async(main_worker, args=(files[idx], args, idx))
        pool.close()
        pool.join()
    else:
        main_worker(files[0], args)

    print("---------------------------------------------------------")
    print("处理完成: {}".format(files))
    cost.toc()
    print("---------------------------------------------------------")


def main_worker(bag, args, process_id=0):
    rosbag_dir = args.rosbag_dir
    data_path = args.data_path

    rosbag_path = os.path.join(rosbag_dir, bag)
    dir_name = os.path.join(data_path, bag.split('.')[0])
    print("开始解析<{}>......".format(rosbag_path))

    bag = rosbag.Bag(rosbag_path)
    #
    lidar_topics = ["/carla/ego_vehicle/lidar1", "/carla/ego_vehicle/lidar2", "/carla/ego_vehicle/lidar3",
                    "/carla/ego_vehicle/lidar4"]
    radar_topics = ["/carla/ego_vehicle/radar3"]
    old_radar_topics = {}
    # camera_topics = ["/camera71/compressed", "/camera72/compressed", "/camera73/compressed", "/camera74/compressed",
    #                  "/camera75/compressed", "/camera76/compressed",
    #                  "/camera77/compressed", "/camera78/compressed", "/camera79/compressed", "/camera80/compressed",
    #                  "/camera81/compressed", "/camera82/compressed"]
    camera_topics = []
    camera_ori_topics = ["/carla/ego_vehicle/camera75/CompressedImage",
                         "/carla/ego_vehicle/camera77/CompressedImage",
                         "/carla/ego_vehicle/camera80/CompressedImage",
                         "/carla/ego_vehicle/camera81/CompressedImage"]

    camera_info_topics = ["/carla/ego_vehicle/camera75/camera_info",
                          "/carla/ego_vehicle/camera77/camera_info",
                          "/carla/ego_vehicle/camera80/camera_info",
                          "/carla/ego_vehicle/camera81/camera_info"]

    camera_prefix_topics = ["/carla/ego_vehicle/camera75",
                            "/carla/ego_vehicle/camera77",
                            "/carla/ego_vehicle/camera80",
                            "/carla/ego_vehicle/camera81"]

    # camera_extra_names = ["SceneColor", "SceneDepth", "GBufferA", "GBufferB", "GBufferC", "GBufferD", "Velocity",
    #                       "GBufferSSAO", "CustomStencil", "Semantic", "SemanticColor"]
    camera_extra_names = ["SceneColor", "SceneDepth", "GBufferA", "GBufferB", "GBufferC", "GBufferD", "Velocity",
                          "GBufferSSAO", "CustomStencil", "Semantic"]
    # lidar_topics = ["/pointcloud_lidar1", "/pointcloud_lidar2", "/pointcloud_lidar3"]
    # radar_topics = ["/pointcloud_radar3"]
    # old_radar_topics = {}
    # camera_topics = []
    # camera_ori_topics = ["/camera71", "/camera72", "/camera73", "/camera74",
    #                  "/camera75", "/camera76",
    #                   "/camera79", "/camera80",
    #                  "/camera82"]

    print("保存在<{}>".format(dir_name))

    # 新建目录
    new_dir = []
    ffarther_dir = []
    farther_dir = []
    child_dir = []
    ffarther_dir.append(dir_name)
    new_dir = new_dir + ffarther_dir
    farther_dir.append(os.path.join(dir_name, "lidar"))
    farther_dir.append(os.path.join(dir_name, "radar"))  # 毫米波所以直接保存为bin:因为1)不直接标注 2)pcd不能保留r_speed和rcs
    farther_dir.append(os.path.join(dir_name, "label"))
    farther_dir.append(os.path.join(dir_name, "camera"))
    farther_dir.append(os.path.join(dir_name, "camera_extra"))
    farther_dir.append(os.path.join(dir_name, "localization"))
    farther_dir.append(os.path.join(dir_name, "calib"))
    new_dir = new_dir + farther_dir
    for topic in lidar_topics:
        child_dir.append(os.path.join(dir_name, "lidar", topic2path(topic, -1)))
    for topic in radar_topics:
        child_dir.append(os.path.join(dir_name, "radar", topic2path(topic, -1)))
    for topic in camera_topics:
        child_dir.append(os.path.join(dir_name, "camera", topic2path(topic)))
    for topic in camera_ori_topics:
        child_dir.append(os.path.join(dir_name, "camera", topic2path(topic, index=-2)))
    for camera in camera_prefix_topics:
        child_dir.append(os.path.join(dir_name, "camera_extra", topic2path(camera)))
        for extra_name in camera_extra_names:
            child_dir.append(os.path.join(dir_name, "camera_extra", topic2path(camera), extra_name))

    new_dir = new_dir + child_dir
    for dir in new_dir:
        if not os.path.exists(dir):
            print("新建目录<{}>".format(dir))
            os.makedirs(dir, exist_ok=True)

    if os.path.exists(os.path.join(rosbag_dir, 'calib.json')):
        copy_file(os.path.join(rosbag_dir, 'calib.json'), os.path.join(dir_name, 'calib'))
        print("拷贝 calib.json , from {} to {}".format(os.path.join(rosbag_dir, 'calib'),
                                                       os.path.join(dir_name, 'calib')))
    else:
        # 创建calib.yml
        eye44 = np.eye(4).tolist()  # 4*4的旋转平移矩阵
        eye33 = np.eye(3).tolist()  # 3*3的旋转平移矩阵
        calib_file = os.path.join(dir_name, 'calib', 'calib.json')
        calib_data = dict()
        calib_data['lidar1'] = dict()
        calib_data['lidar1']['transforms'] = dict()
        calib_data['lidar1']['transforms']['base_link'] = np.array(
            [[0.000796064, 0.9998, -0.0199804, 0],
             [-0.999735, 0.000336186, -0.0230093, -1.4],
             [-0.022998, 0.0199934, 0.999536, 1.28],
             [0.0, 0.0, 0.0, 1.0]]).tolist()
        calib_data['lidar2'] = dict()
        calib_data['lidar2']['transforms'] = dict()
        calib_data['lidar2']['transforms']['base_link'] = np.array(
            [[0.0107962, -0.999492, -0.0299938, 0],
             [0.999942, 0.0107913, 0.000323837, 1.3],
             [0, -0.0299955, 0.99955, 1.35],
             [0.0, 0.0, 0.0, 1.0]]).tolist()
        calib_data['lidar3'] = dict()
        calib_data['lidar3']['transforms'] = dict()
        calib_data['lidar3']['transforms']['base_link'] = np.array(
            [[0.99995, 9.99967e-05, 0.00999933, 2],
             [0, 0.99995, -0.00999983, 0],
             [-0.00999983, 0.00999933, 0.9999, 1.1],
             [0.0, 0.0, 0.0, 1.0]]).tolist()
        calib_data['lidar4'] = dict()
        calib_data['lidar4']['transforms'] = dict()
        calib_data['lidar4']['transforms']['base_link'] = np.array(
            [[-0.999944, -0.00369245, -0.00996334, -3.1],
             [0.00359245, -0.999943, 0.0100357, 0.2],
             [-0.00999983, 0.00999933, 0.9999, 1.3],
             [0.0, 0.0, 0.0, 1.0]]).tolist()
        calib_data['radar3'] = dict()
        calib_data['radar3']['transforms'] = dict()
        calib_data['radar3']['transforms']['base_link'] = np.array(
            [[1.0, 0.0, 0.0, 1.8],
             [0.0, 1.0, 0.0, 0.0],
             [0.0, 0.0, 1.0, 1.6],
             [0.0, 0.0, 0.0, 1.0]]).tolist()
        for cam_id in ['camera75', 'camera77', 'camera80', 'camera81']:
            calib_data[cam_id] = dict()
            calib_data[cam_id]['transforms'] = dict()
        calib_data['camera75']['transforms']['base_link'] = np.array(
            [[0.003022734951638418, -0.0013826472978586283, 0.9999949513105737, 2.0552528163239945],
             [-0.9999264042966791, 0.011762464169231247, 0.0030387991432021906, 0.21210449149924276],
             [-0.011766637379637146, -0.9999297273788396, -0.0013469934340480324, 1.0876293178270175],
             [0.0, 0.0, 0.0, 1.0]]).tolist()
        calib_data['camera77']['transforms']['base_link'] = np.array(
            [[-0.01261539065140769, 0.025386277886985834, -0.9995982266215556, -3.2461425468147542],
             [0.9995891542880018, -0.025393770969560678, -0.013260201505974819, -0.06747394112787104],
             [-0.025720214696331622, -0.9993552761455446, -0.025055534375152348, 1.4508716834658448],
             [0.0, 0.0, 0.0, 1.0]]).tolist()
        calib_data['camera80']['transforms']['base_link'] = np.array(
            [[-0.9995041250105752, -0.009716197150941653, 0.029970387802864533, 0.06404694658754413],
             [-0.029860843168041655, -0.01123280625705582, -0.9994908963344377, -1.4240985372582626],
             [0.010047858147755731, -0.9998894079398011, 0.010937087636534688, 1.268783890537187],
             [0.0, 0.0, 0.0, 1.0]]).tolist()
        calib_data['camera81']['transforms']['base_link'] = np.array(
            [[0.9997917334074543, 0.007837232625340699, -0.018836124062913726, -0.13316415534134804],
             [0.019239867752761645, -0.055091549688316525, 0.9982958432905688, 1.485195144098181],
             [0.006786164435152244, -0.9984500953057992, -0.05523089115557751, 1.3635967951903432],
             [0.0, 0.0, 0.0, 1.0]]).tolist()

        # 写入JSON文件
        with open(calib_file, 'w', encoding='utf-8') as file:
            json.dump(calib_data, file, ensure_ascii=False, indent=2)

        # fs = cv2.FileStorage(os.path.join(dir_name, 'calib', 'sensor.yml'), cv2.FileStorage_WRITE)
        # matrix = np.eye(4)  # 4*4的旋转平移矩阵
        # for topic in lidar_topics + radar_topics:
        #     sensor_id = topic[-6:]
        #     fs.write(sensor_id, matrix)
        # fs.release()
        print("新建标定文件 in {}".format(calib_file))

    # 创建或打开文件
    csvfile = open(os.path.join(dir_name, "localization", 'localization.csv'), mode='w', newline='')
    # 标题列表
    fieldnames = ['stamp(ms)', 'latitude', 'longitude', 'altitude', 'mercator_x', 'mercator_y', 'mercator_z',
                  'vel_x', 'vel_y', 'vel_z', 'accel_x', 'accel_y', 'accel_z',
                  'roll', 'pitch', 'yaw', 'angular_x', 'angular_y', 'angular_z']
    # 创建 DictWriter 对象
    write = csv.DictWriter(csvfile, fieldnames=fieldnames)
    # 写入表头
    write.writeheader()

    # 自车id
    ego_vehicle_id = -1
    for topic, msg, t in bag:
        if topic in '/carla/ego_vehicle/vehicle_info':
            ego_vehicle_id = msg.id

        if ego_vehicle_id == -1:
            continue

        for extra_name in camera_extra_names:
            if extra_name in topic:
                np_arr = np.frombuffer(msg.data, dtype=np.uint8)  # 将消息中的字节数据转换为 numpy 数组
                cv_image = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
                cv2.imwrite(
                    os.path.join(dir_name, "camera_extra", topic2path(topic, -2), extra_name,
                                 str(get_ms(msg, process_id)) + ".jpg"),
                    cv_image)
                continue

        if topic in lidar_topics:
            lidar = pc2.read_points(msg)
            points = np.array(list(lidar))
            # # 创建一个PointCloud对象
            # pcd = o3d.geometry.PointCloud()
            # # 将随机数转换成PointCloud点数据
            # pcd.points = o3d.utility.Vector3dVector(points[:, :3])
            # # 将PointCloud点数据保存成pcd文件，格式为assii文本格式
            # o3d.io.write_point_cloud(os.path.join(dir_name,'lidar', topic2path(topic), str(get_ms(msg, process_id)) + ".pcd"), pcd,
            #                          write_ascii=True)
            assert points.shape[1] >= 4
            write_points(points[:, [0, 1, 2, 3]],
                         os.path.join(dir_name, 'lidar', topic2path(topic, -1), str(get_ms(msg, process_id)) + ".bin"))

        elif topic in radar_topics:
            radar = pc2.read_points(msg)
            points = np.array(list(radar))
            points = points[:, :7]  # x y z r_speed rcs
            file_path = os.path.join(dir_name, 'radar', topic2path(topic, -1), str(get_ms(msg, process_id)) + ".bin")
            write_points(points, file_path)
            # # 创建一个PointCloud对象
            # pcd = o3d.geometry.PointCloud()
            # # 将随机数转换成PointCloud点数据
            # pcd.points = o3d.utility.Vector3dVector(points[:, :3])# x y z r_speed rcs
            # # 将PointCloud点数据保存成pcd文件，格式为assii文本格式
            # o3d.io.write_point_cloud(os.path.join(dir_name,'radar', topic2path(topic), str(get_ms(msg, process_id)) + ".pcd"), pcd,
            #                          write_ascii=True)


        elif topic in camera_topics:
            bridge = CvBridge()
            cv_image = bridge.compressed_imgmsg_to_cv2(msg, "bgr8")
            if topic[:9] == '/camera81' or topic[:9] == '/camera80':  # 前后相机超前600ms，左右相机超前500ms
                cv2.imwrite(
                    os.path.join(dir_name, "camera", topic2path(topic), str(get_ms(msg, process_id) + 450) + ".jpg"),
                    cv_image)
            else:
                cv2.imwrite(
                    os.path.join(dir_name, "camera", topic2path(topic), str(get_ms(msg, process_id) + 600) + ".jpg"),
                    cv_image)
            # cv2.imshow("Image window", cv_image)
            # cv2.waitKey(3)
        elif topic in camera_ori_topics:
            np_arr = np.frombuffer(msg.data, dtype=np.uint8)  # 将消息中的字节数据转换为 numpy 数组
            cv_image = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
            # cv2.imwrite(os.path.join(dir_name, "camera", topic2path(topic), str(get_ms(msg, process_id)+600) + ".jpg"), cv_image)
            cv2.imwrite(
                os.path.join(dir_name, "camera", topic2path(topic, index=-2), str(get_ms(msg, process_id)) + ".jpg"),
                cv_image)

        elif topic in camera_info_topics:
            # 写入内参和尺度
            cam_id = topic.split('/')[-2]
            if 'ori' not in calib_data[cam_id]:
                calib_data[cam_id]['ori'] = dict()
                calib_data[cam_id]['ori']['imgh'] = msg.height
                calib_data[cam_id]['ori']['imgw'] = msg.width
                calib_data[cam_id]['ori']['K'] = np.array(msg.K).reshape((3, 3)).tolist()
                calib_data[cam_id]['ori']['D'] = np.zeros(shape=(5)).tolist()
                with open(calib_file, 'w', encoding='utf-8') as file:
                    json.dump(calib_data, file, ensure_ascii=False, indent=2)

        elif topic == '/carla/ego_vehicle/odometry':
            roll, pitch, yaw = quaternion_to_euler_angles(
                [msg.pose.pose.orientation.w, msg.pose.pose.orientation.x, msg.pose.pose.orientation.y,
                 msg.pose.pose.orientation.z], axes='sxyz')
            write.writerow(
                {'stamp(ms)': str(get_ms(msg, process_id)) + '\t',
                 'latitude': 0.0, 'longitude': 0.0, 'altitude': 0.0,
                 'mercator_x': msg.pose.pose.position.x,
                 'mercator_y': msg.pose.pose.position.y,
                 'mercator_z': msg.pose.pose.position.z,
                 'vel_x': 0.0, 'vel_y': 0.0, 'vel_z': 0.0,
                 'accel_x': 0.0, 'accel_y': 0.0, 'accel_z': 0.0,
                 'roll': roll, 'pitch': pitch, 'yaw': yaw,
                 'angular_x': 0.0, 'angular_y': 0.0, 'angular_z': 0.0})
        elif topic == '/carla/objects':
            ego = None
            for obj in msg.objects:
                if obj.id == ego_vehicle_id:
                    ego = obj
            assert ego is not None

            object_datas = []
            for obj in msg.objects:
                from cyw_devkit.core.zmath import relative_pose
                if obj.id == ego_vehicle_id:
                    # roll, pitch, yaw = quaternion_to_euler_angles([obj.pose.orientation.w, obj.pose.orientation.x, obj.pose.orientation.y,
                    #                            obj.pose.orientation.z], axes='sxyz')
                    # write.writerow(
                    #     {'stamp(ms)': str(get_ms(msg, process_id)) + '\t',
                    #      'latitude': 0.0, 'longitude': 0.0, 'altitude': 0.0,
                    #      'mercator_x': obj.pose.position.x,
                    #      'mercator_y': obj.pose.position.y,
                    #      'mercator_z': obj.pose.position.z,
                    #      'vel_x': 0.0, 'vel_y': 0.0, 'vel_z': 0.0,
                    #      'accel_x': 0.0, 'accel_y': 0.0, 'accel_z': 0.0,
                    #      'roll':roll, 'pitch':pitch, 'yaw': yaw,
                    #      'angular_x': 0.0, 'angular_y': 0.0, 'angular_z': 0.0})
                    continue  # 不写自己的box

                relative_pose_info = relative_pose(ego.pose, obj.pose)
                assert int(obj.classification) in idx2cls.keys()

                obj_type = idx2cls[int(obj.classification)]
                if obj.shape.dimensions[0] > 10:
                    print('bus 映射错误，去carla工程修复')
                    obj_type = 'Bus'
                distance = relative_pose_info['dx'] ** 2 + relative_pose_info['dy'] ** 2
                if distance > cls2dis[obj_type] * cls2dis[obj_type]:
                    continue
                z = relative_pose_info['dz']
                if obj_type != 'Pedestrian':  # 车的z定义似乎是底部高度
                    z += 0.5 * obj.shape.dimensions[2]
                object_data = {
                    'obj_id': int(obj.id),
                    'obj_type': obj_type,
                    'psr': {
                        'position': {
                            'x': relative_pose_info['dx'],
                            'y': relative_pose_info['dy'],
                            'z': z
                        },
                        'rotation': {
                            'x': 0.0,
                            'y': 0.0,
                            'z': relative_pose_info['dyaw']
                        },
                        'scale': {
                            'x': obj.shape.dimensions[0],
                            'y': obj.shape.dimensions[1],
                            'z': obj.shape.dimensions[2]
                        }
                    }
                }
                object_datas.append(object_data)
            with open(os.path.join(dir_name, 'label', str(get_ms(msg, process_id)) + '.json'), 'w',
                      encoding='utf-8') as file:
                json.dump(object_datas, file, ensure_ascii=False, indent=2)

    for dir in child_dir:
        print("文件夹:<{}> 包含文件<{}>个".format(dir.split('/')[-1], len(os.listdir(dir))))

    bag.close()


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Configuration Parameters')
    parser.add_argument('--rosbag-dir',
                        default='/home/adt/bags/work_space/bags',
                        help='rosbag')
    parser.add_argument('--data-path',
                        default='/home/adt/bags/work_space/datasets')
    args = parser.parse_args()

    main(args)
